//===- ScalarEvolutionTestPass.cpp - Test Pass using Scalar Evolution -----===//
//
//                     Damsl Lab LLVM-MPI Project
//
//
//===----------------------------------------------------------------------===//
//
//
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "arraytracking"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/ValueMap.h"
#include "llvm/Analysis/DebugInfo.h"
#include "llvm/Analysis/Passes.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/MPI/MPI.h"
#include "llvm/Pass.h"
#include "llvm/Constant.h"
#include "llvm/Module.h"
#include "llvm/Instructions.h"
#include <vector>
#include <map>
#include <algorithm>
#include <string>
#include <cctype>   //enable char.isdigit()
using namespace llvm;

namespace {
// scalar value is one element array
class ArraySrcTracking : public FunctionPass {
  ScalarEvolution *SE;
  // store result
  typedef std::map<const Value *, SmallSet<Value*, 6>*> ValueSetMapTy;
  ValueSetMapTy VMap;

public:
  static char ID; // Class identification, replacement for typeinfo
  ArraySrcTracking() : FunctionPass(ID), SE(0) {}

private:
  Value* traceup(Value* v);
  Value* traceup_out_of_procudure(Argument* arg);
  void tracedown(Value* v, SmallSet<Value*, 6> &result);
  bool isMpiCommCall(const Instruction &II);
  bool isMpiSendCall(const Instruction &II);
  inline bool funcname(const CallInst* II, const char* name) {
    Function* func = II->getCalledFunction();
    StringRef fname = func->getName();
    if (fname.startswith(StringRef(name))) {
      return true;
    } else {return false;}
  }
  inline bool funcname(const CallInst* II, StringRef name) {
    Function* func = II->getCalledFunction();
    StringRef fname = func->getName();
    if (fname.startswith(name)) {
      return true;
    } else {return false;}
  }
  void getAnalysisUsage(AnalysisUsage &AU) const;
  bool runOnFunction(Function &F);
  ~ArraySrcTracking();
};

// Register this pass...
char ArraySrcTracking::ID = 0;
static RegisterPass<ArraySrcTracking> X("arraytracking",
                              "track src of array", false, true);

}  // End of anonymous namespace


void
ArraySrcTracking::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.addRequiredTransitive<ScalarEvolution>();
  //AU.addRequired<PromotePass>();
  AU.setPreservesAll();
}

// FIXME: return the first source.
// (improve this) there could be multiple sources.
Value* ArraySrcTracking::traceup_out_of_procudure(Argument* arg) {
  DEBUG(dbgs() << "begin traceup_out_of_procudure\n");
  // current function
  Function* F = arg->getParent();
  // current module
  Module* M = F->getParent();
  int argindex = arg->getArgNo();

  for(Module::iterator I = M->begin(), E = M->end(); I != E; ++I) {
    for (Function::const_iterator FI = I->begin(), FE = I->end();
        FI != FE; ++FI)
      for (BasicBlock::const_iterator BI = FI->begin(), BE = FI->end();
          BI != BE; ++BI) {
        if(const CallInst* callInst = dyn_cast<CallInst>(&*BI)) {
          if(funcname(callInst, F->getName())) {
            return traceup(callInst->getArgOperand(argindex));
          }
        }

      }
  }

  return arg;
}

// see if array is allocated in this function
Value* ArraySrcTracking::traceup(Value* v) {
  // v is an instruction
  //  if(Instruction* inst = dyn_cast<Instruction>(v)) {
  //    // v's parent basicblock
  //    BasicBlock* par = inst->getParent();
  //    // v's parent function
  //    Function* fun = inst->getParent()->getParent();
  //    // v's parent basicblock is a entry block?
  //    if(par == fun->getEntryBlock()) {
  //      // v is alloca in entry block. call by value
  //      if(AllocaInst* allo = dyn_cast<AllocaInst>(v))
  //        return v;
  //    }
  //  }
  DEBUG(dbgs() << "begin traceup\n");
  //  if(v->hasName()) {
  //    DEBUG(dbgs() << "hasName\n");
  //  } else {
  //    DEBUG(dbgs() << "not hasName\n");
  //  }
  DEBUG(dbgs() << "value name is " << v->getNameStr() << "\n");


  // v is not an instruction and has a name.
  // v maybe is a array or it is a index.
  Instruction* inst = dyn_cast<Instruction>(v);

  if(!inst) {
    if(v->hasName()) {


      // if v is argument, trace out of current function
      if(Argument* arg = dyn_cast<Argument>(v)) {
        DEBUG(dbgs() << v->getNameStr() << " is argument of a function\n");
        v = traceup_out_of_procudure(arg);
      }

      //DEBUG(dbgs() << v->getNameStr() << " is not instruction and has name\n");
      DEBUG(dbgs() << "!!!found. value name is " << v->getNameStr() << "\n");
      return v;
    } else {
      DEBUG(dbgs() << "wrong: v is not instruction and has no name.\n");
    }
  }

  // check if vector7 = bitcast vector ....
  bool sameprefix = false;
  if(v->hasName()) {
    DEBUG(dbgs() << "in sameprefix\n");
    for (User::op_iterator i = inst->op_begin(),
        e = inst->op_end(); i != e; ++i) {
      if((*i)->hasName() && v->getName().startswith((*i)->getName())) {
        DEBUG(dbgs() << "sameprefix is true\n");
        sameprefix = true;
        break;
      }
    }

    if(sameprefix == false){
      DEBUG(dbgs() << "found array def. sameprfix false.\n");
      DEBUG(dbgs() << "!!!found. value name is " << v->getNameStr() << "\n");
      return v;
    }
  }

  //DEBUG(dbgs() << "after dyncast\n");

  //  Value *t;
  //  // find instruction to trace up to.
  //  for (User::op_iterator i = inst->op_begin(),
  //      e = inst->op_end(); i != e; ++i) {
  //    t = *i;
  //    if(dyn_cast<Instruction>(t)) {
  //      break;
  //    }
  //  }
  //
  //  DEBUG(dbgs() << "after for loop\n");
  //
  //  if(t == NULL) {
  //    // TODO
  //    DEBUG(dbgs() << "v has no op that is instruction. what to do?\n");
  //    return t; // return null
  //  }

  User::op_iterator i = inst->op_begin();
  Value *t = *i;
  DEBUG(dbgs() << "before iterate\n");
  return traceup(t);
}

// input v: array; output result: instructions that define v;
void ArraySrcTracking::tracedown(Value* v, SmallSet<Value*, 6> &result) {
  // store instruction define array
  // TODO: if first op of store instruction is function parameter. array
  // is defined outside of this function. need interprocedural analysis.
  DEBUG(dbgs() << "begin tracedown\n");

  for (Value::use_iterator i = v->use_begin(), e = v->use_end(); i != e; ++i) {
    User *U = *i;
    //    DEBUG(dbgs() << v->getNumUses() << "\n");
    //
    //
    //    if(BitCastInst* inst = dyn_cast<BitCastInst>(U)) {
    //      DEBUG(dbgs() << "is BitCastInst\n");
    //    }
    //
    if(LoadInst* inst = dyn_cast<LoadInst>(U)) {
      DEBUG(dbgs() << "is LoadInst, return.\n");
      continue;
    }

    //DEBUG(dbgs() << "before StoreInst.\n");
    if(StoreInst* inst = dyn_cast<StoreInst>(U)) {
      DEBUG(dbgs() << "is StoreInst\n");
      if(inst->getPointerOperand() == v) {
      	DEBUG(dbgs() << "    ***StoreInst\n");
        result.insert(*i);
        return;
      }
    }
    //DEBUG(dbgs() << "after StoreInst.\n");

    //DEBUG(dbgs() << "before CallInst.\n");
    // If user is function call instruction
    if (const CallInst* callInst = dyn_cast<CallInst>(U)) {
      // printf won't define array
      if(funcname(callInst, "printf")) {
        continue;
      }

      if (!isMpiSendCall(*callInst)) {
      	DEBUG(dbgs() << "    ***CallInst. function name is "
               << callInst->getCalledFunction()->getName() << "\n");
        result.insert(*i);
        return;
      }
    }
    //DEBUG(dbgs() << "after CallInst.\n");

    tracedown(*i, result);
  }

  return;
}

bool ArraySrcTracking::runOnFunction(Function &F) {
  SE = &getAnalysis<ScalarEvolution>();
  DEBUG(dbgs() << "\n\nFunction: " << F.getName() << "\n");
  // search for MPI send
  for (Function::iterator BI = F.begin(), BE = F.end(); BI
  != BE; ++BI) {
    for (BasicBlock::iterator II = (*BI).begin() , E = (*BI).end(); II
    != E; ++II) {
      //DEBUG(dbgs() << "iteration\n");
      if (isMpiSendCall(*II)) {
        //DEBUG(dbgs() << "  isMpiSendCall\n");
        CallInst* sendcall = dyn_cast<CallInst>(&*II);
        //DEBUG(dbgs() << "    " << sendcall->getNumOperands() << "\n");
        Value* array_ptr = sendcall->getOperand(0);
        DEBUG(dbgs() << "runOnFunction: Traceup begin.\n");
        Value* src = traceup(array_ptr);
        //DEBUG(dbgs() << "     after traceup\n");
        SmallSet<Value*, 6>* t = new SmallSet<Value*, 6>;
        DEBUG(dbgs() << "runOnFunction: Tracedown begin.\n");
        DEBUG(dbgs() << "  value's name: " << src->getNameStr() << "\n");
        tracedown(src, *t);
        VMap[array_ptr] = t;
      }
    }
  }
  return false;
}

ArraySrcTracking::~ArraySrcTracking() {
  for(ValueSetMapTy::const_iterator i = VMap.begin(), e = VMap.end();
      i != e; i++) {
    delete std::pair<const Value *, SmallSet<Value*, 6>*>(*i).second;
  }
}

// Return true if Instruction II is a MPI Communication call.
bool ArraySrcTracking::isMpiCommCall(const Instruction &II) {
  if (const CallInst* callInst = dyn_cast<CallInst>(&II)) {
    Function* func = callInst->getCalledFunction();
    StringRef fname = func->getName();
    if (fname.startswith(StringRef("MPI_Send"))  ||
        fname.startswith(StringRef("MPI_Isend")) ||
        fname.startswith(StringRef("MPI_Recv"))  ||
        fname.startswith(StringRef("MPI_Irecv")))
      return true;
  }
  return false;
}

// Return true if Instruction II is a MPI send call.
bool ArraySrcTracking::isMpiSendCall(const Instruction &II) {
  if (const CallInst* callInst = dyn_cast<CallInst>(&II)) {
    Function* func = callInst->getCalledFunction();
    StringRef fname = func->getName();
    if (fname.startswith(StringRef("MPI_Send"))  ||
        fname.startswith(StringRef("MPI_Isend")))
      return true;
  }
  return false;
}


